{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import open3d as o3d\n",
    "import numpy as np\n",
    "import config\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Open3D WARNING] GLFW Error: WGL: Failed to make context current: The handle is invalid. \n"
     ]
    }
   ],
   "source": [
    "pcl = o3d.io.read_point_cloud(\"pointcloud.ply\")\n",
    "o3d.visualization.draw_geometries([pcl])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remove the ground plane"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Open3D WARNING] GLFW Error: WGL: Failed to make context current: The requested transformation operation is not supported. \n",
      "[Open3D WARNING] GLFW Error: WGL: Failed to make context current: The requested transformation operation is not supported. \n",
      "[Open3D WARNING] GLFW Error: WGL: Failed to make context current: The handle is invalid. \n",
      "[Open3D WARNING] GLFW Error: WGL: Failed to make context current: The requested transformation operation is not supported. \n",
      "[Open3D WARNING] GLFW Error: WGL: Failed to make context current: The requested transformation operation is not supported. \n"
     ]
    }
   ],
   "source": [
    "points = o3d.utility.Vector3dVector(np.array([\n",
    "\t[config.point_cloud_range[\"x_min\"], config.point_cloud_range[\"y_min\"], config.min_box_height],\n",
    "\t[config.point_cloud_range[\"x_max\"], config.point_cloud_range[\"y_max\"], config.point_cloud_range[\"z_max\"]]\n",
    "]))\n",
    "box_pcl = pcl.crop(\n",
    "\tbounding_box = o3d.geometry.AxisAlignedBoundingBox.create_from_points(points)\n",
    ")\n",
    "o3d.visualization.draw_geometries([box_pcl])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_plane_vec_constraint(norm_vec, pts, thresh=0.05, n_iterations=300):\n",
    "\tbest_eq = []\n",
    "\tbest_inliers = []\n",
    "\n",
    "\tn_points = pts.shape[0]\n",
    "\tfor iter in range(n_iterations):\n",
    "\t\tid_sample = random.sample(range(0, n_points), 1)\n",
    "\t\tpoint = pts[id_sample]\n",
    "\t\td = -np.sum(np.multiply(norm_vec, point))\n",
    "\t\tplane_eq = [*norm_vec, d]\n",
    "\t\tpt_id_inliers = get_plane_inliers(plane_eq, pts, thresh)\n",
    "\t\tif len(pt_id_inliers) > len(best_inliers):\n",
    "\t\t\tbest_eq = plane_eq\n",
    "\t\t\tbest_inliers = pt_id_inliers\n",
    "\n",
    "\treturn best_eq, best_inliers\n",
    "\n",
    "def get_plane_inliers(plane_eq, pts, thresh=0.05):\n",
    "\tpt_id_inliers = []\n",
    "\tdist_pt = get_pts_distances_plane(plane_eq, pts)\n",
    "\n",
    "\t# Select indexes where distance is bigger than the threshold\n",
    "\tpt_id_inliers = np.where(np.abs(dist_pt) <= thresh)[0]\n",
    "\treturn pt_id_inliers\n",
    "\n",
    "def get_pts_distances_plane(plane_eq, pts):\n",
    "\tdist_pt = (plane_eq[0] * pts[:, 0] + plane_eq[1] * pts[:, 1] \n",
    "\t\t\t+ plane_eq[2] * pts[:, 2] + plane_eq[3])\\\n",
    "\t\t\t/ np.sqrt(plane_eq[0] ** 2 + plane_eq[1] ** 2 + plane_eq[2] ** 2)\n",
    "\treturn dist_pt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remove the top plane"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "points_np = np.asarray(box_pcl.points)\n",
    "top_plane_eq, top_plane_inliers = fit_plane_vec_constraint([0, 0, 1], points_np, 0.01, 30)\n",
    "\n",
    "top_plane = box_pcl.select_by_index(top_plane_inliers)\n",
    "top_side_pcl = top_plane\n",
    "height = abs(top_plane_eq[3])\n",
    "\n",
    "points = o3d.utility.Vector3dVector(np.array([\n",
    "\t[config.point_cloud_range[\"x_min\"], config.point_cloud_range[\"y_min\"], config.point_cloud_range[\"z_min\"]],\n",
    "\t[config.point_cloud_range[\"x_max\"], config.point_cloud_range[\"y_max\"], height - config.min_box_height]\n",
    "]))\n",
    "box_sides_pcl = box_pcl.crop(\n",
    "\tbounding_box = o3d.geometry.AxisAlignedBoundingBox.create_from_points(points)\n",
    ")\n",
    "\n",
    "o3d.visualization.draw_geometries([box_sides_pcl])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit_vertical_plane(norm_vec, pts, thresh=0.05, n_iterations=300):\n",
    "\tbest_eq = []\n",
    "\tbest_inliers = []\n",
    "\n",
    "\tn_points = pts.shape[0]\n",
    "\tfor iter in range(n_iterations):\n",
    "\t\tid_sample = random.sample(range(0, n_points), 1)\n",
    "\t\tpoint = pts[id_sample]\n",
    "\t\td = -np.sum(np.multiply(norm_vec, point))\n",
    "\t\tplane_eq = [*norm_vec, d]\n",
    "\t\tpt_id_inliers = get_plane_inliers(plane_eq, pts, thresh)\n",
    "\t\tif len(pt_id_inliers) > len(best_inliers):\n",
    "\t\t\tbest_eq = plane_eq\n",
    "\t\t\tbest_inliers = pt_id_inliers\n",
    "\n",
    "\treturn best_eq, best_inliers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Open3D WARNING] GLFW Error: WGL: Failed to make context current: The requested transformation operation is not supported. \n"
     ]
    }
   ],
   "source": [
    "plane_eq, plane_inliers = box_sides_pcl.segment_plane(distance_threshold=0.01, ransac_n=3, num_iterations=100)\n",
    "side_pcl = box_sides_pcl.select_by_index(plane_inliers)\n",
    "side_pcl.paint_uniform_color([1, 0, 0])\n",
    "\n",
    "o3d.visualization.draw_geometries([side_pcl, box_sides_pcl])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.7 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.7"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "570feb405e2e27c949193ac68f46852414290d515b0ba6e5d90d076ed2284471"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
